import sys
import os
import numpy as np
import pandas as pd
from PyQt5.QtWidgets import (
    QApplication, QWidget, QVBoxLayout, QLabel, QPushButton,
    QFileDialog, QLineEdit, QHBoxLayout, QMessageBox
)
import plotly.graph_objs as go
from plotly.offline import plot
import tempfile

class WellTrajectoryApp(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("三维井眼轨迹可视化工具")
        self.resize(1200, 800)

        self.well_data = {}  # 存储多口井的数据

        self.layout = QVBoxLayout()
        self.input_layout = QHBoxLayout()

        self.well_name_input = QLineEdit()
        self.well_name_input.setPlaceholderText("井号")
        self.x_input = QLineEdit()
        self.x_input.setPlaceholderText("X坐标")
        self.y_input = QLineEdit()
        self.y_input.setPlaceholderText("Y坐标")

        self.input_layout.addWidget(self.well_name_input)
        self.input_layout.addWidget(self.x_input)
        self.input_layout.addWidget(self.y_input)

        self.upload_btn = QPushButton("上传Excel文件")
        self.upload_btn.clicked.connect(self.load_excel)

        self.download_template_btn = QPushButton("下载Excel模板")
        self.download_template_btn.clicked.connect(self.download_template)

        self.plot_btn = QPushButton("绘制三维轨迹图")
        self.plot_btn.clicked.connect(self.plot_trajectories)

        self.layout.addLayout(self.input_layout)
        self.layout.addWidget(self.upload_btn)
        self.layout.addWidget(self.download_template_btn)
        self.layout.addWidget(self.plot_btn)
        self.setLayout(self.layout)

    def load_excel(self):
        file_path, _ = QFileDialog.getOpenFileName(self, "选择Excel文件", "", "Excel Files (*.xlsx *.xls)")
        if not file_path:
            return
        try:
            df = pd.read_excel(file_path)
            assert set(['MD', 'INCL', 'AZI']).issubset(df.columns)

            md = df['MD'].to_numpy()  # 测深
            inc = np.radians(df['INCL'].to_numpy())  # 井斜角
            azi = np.radians(df['AZI'].to_numpy())  # 方位角

            x0 = float(self.x_input.text())
            y0 = float(self.y_input.text())

            x, y, z = [x0], [y0], [0.0]

            # 最小曲率法计算每段坐标增量
            for i in range(1, len(md)):
                md1, md2 = md[i-1], md[i]
                inc1, inc2 = inc[i-1], inc[i]
                azi1, azi2 = azi[i-1], azi[i]
                d_md = md2 - md1

                # 计算狗腿度 (dogleg angle)
                cos_dl = np.cos(inc2 - inc1) - np.sin(inc1) * np.sin(inc2) * (1 - np.cos(azi2 - azi1))
                # 数值保护
                cos_dl = np.clip(cos_dl, -1.0, 1.0)
                dl = np.arccos(cos_dl)

                # 计算比率因子 RF
                if dl > 1e-6:
                    rf = (2 / dl) * np.tan(dl / 2)
                else:
                    rf = 1.0

                # 计算坐标增量
                dx = (d_md / 2) * (np.sin(inc1) * np.sin(azi1) + np.sin(inc2) * np.sin(azi2)) * rf
                dy = (d_md / 2) * (np.sin(inc1) * np.cos(azi1) + np.sin(inc2) * np.cos(azi2)) * rf
                dz = (d_md / 2) * (np.cos(inc1) + np.cos(inc2)) * rf

                x.append(x[-1] + dx)
                y.append(y[-1] + dy)
                z.append(z[-1] + dz)

            self.well_data[self.well_name_input.text()] = {
                'x': x,
                'y': y,
                'z': z
            }
            QMessageBox.information(self, "成功", f"成功加载井：{self.well_name_input.text()}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"加载文件失败：{str(e)}")

    def download_template(self):
        template_df = pd.DataFrame({
            'MD': [0, 100, 200],
            'INCL': [0, 10, 20],
            'AZI': [0, 90, 180]
        })
        file_path, _ = QFileDialog.getSaveFileName(self, "保存模板", "well_template.xlsx", "Excel Files (*.xlsx)")
        if file_path:
            template_df.to_excel(file_path, index=False)
            QMessageBox.information(self, "下载成功", f"模板已保存至 {file_path}")

    def plot_trajectories(self):
        if not self.well_data:
            QMessageBox.warning(self, "警告", "请先上传至少一个井的数据")
            return
        fig = go.Figure()
        for name, data in self.well_data.items():
            fig.add_trace(go.Scatter3d(
                x=data['x'], y=data['y'], z=data['z'],
                mode='lines+markers', name=name,
                line=dict(width=4), marker=dict(size=3)
            ))
        # 碰撞预警
        wells = list(self.well_data.items())
        for i in range(len(wells)):
            for j in range(i+1, len(wells)):
                xi, yi, zi = wells[i][1]['x'], wells[i][1]['y'], wells[i][1]['z']
                xj, yj, zj = wells[j][1]['x'], wells[j][1]['y'], wells[j][1]['z']
                dmin = min(np.sqrt((np.array(xi)-np.array(xj))**2 +
                                   (np.array(yi)-np.array(yj))**2 +
                                   (np.array(zi)-np.array(zj))**2))
                if dmin < 5:
                    QMessageBox.warning(self, "碰撞预警", f"井 {wells[i][0]} 与 {wells[j][0]} 距离过近: {dmin:.2f}m")
        fig.update_layout(
            title="三维井眼轨迹图",
            scene=dict(
                xaxis_title='X (m)', yaxis_title='Y (m)',
                zaxis_title='TVD (m)', zaxis=dict(autorange='reversed')
            )
        )
        html_file = tempfile.NamedTemporaryFile(delete=False, suffix='.html')
        plot(fig, filename=html_file.name, auto_open=True)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = WellTrajectoryApp()
    window.show()
    sys.exit(app.exec_())
